#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
@author:  runyuanye
@contact: runyuanye@163.com
"""

import argparse
import os
import sys
import numpy as np

sys.path.append('.')

"""
   对merge的人脸ID进行过滤
"""

MIN_PITCH1 = -18.0
MAX_PITCH1 = 16.0
MIN_YAW1 = -16.0
MAX_YAW1 = 16.0
MIN_ROLL1 = -16.0
MAX_ROLL1 = 16.0

MIN_PITCH2 = -20.0
MAX_PITCH2 = 18.0
MIN_YAW2 = -18.0
MAX_YAW2 = 18.0
MIN_ROLL2 = -18.0
MAX_ROLL2 = 18.0


def read_features(file_list, feature_file, feat_dim):
    faceid_subfaceid_file_features = {}
    feat_file = open(feature_file, 'rb')
    file_count = 0
    with open(file_list, 'r') as file:
        for line in file.readlines():
            file_count += 1
            if file_count % 10000 == 0:
                print('Read Feature Count: {}'.format(file_count))

            # if file_count == 10000:
            #     break

            # 20f10478-18a0-414c-a695-c43fb51077aa 11aa6f7c-b776-4e5b-8a16-fac29fbd5cb1 81c996ed4a8a743e046fa09025d69231.jpg 0
            lines = line.strip().split()

            faceid = lines[0]
            subfaceid = lines[1]
            file_name = lines[2]
            # mode = int(lines[3])

            faceid_subfaceid = '{} {}'.format(faceid, subfaceid)

            feature = feat_file.read(feat_dim * 4)

            file_features = faceid_subfaceid_file_features.get(faceid_subfaceid, None)
            if file_features is None:
                files = [line]
                features = [feature]
                file_names = [file_name]
                faceid_subfaceid_file_features[faceid_subfaceid] = [files, features, file_names]
            else:
                file_features[0].append(line)
                file_features[1].append(feature)
                file_features[2].append(file_name)

        print('Read Feature Count: {}'.format(file_count))
    feat_file.close()
    return faceid_subfaceid_file_features


def read_attr_file(file_list):
    file_count = 0
    name_attrs = {}
    with open(file_list, 'r') as file:
        for line in file.readlines():
            file_count += 1
            if file_count % 10000 == 0:
                print('Read Attr Count: {}'.format(file_count))

            # if file_count == 10000:
            #     break

            # gd/zs/nqjd/dxjyjy/150100414a54443452064d2742c93600/video/20190817_082812/000359_000.jpg;-12.010,20.456,50.344;0.630
            lines = line.strip().split(';')
            name = lines[0]
            headpose = lines[1].split(',')
            headpose = list(map(float, headpose))  # pitch, roll, yaw
            quality = float(lines[2])
            name_attrs[name] = [headpose, quality]
        print('Read Attr Count: {}'.format(file_count))
    return name_attrs


def read_face_list(face_list_path):
    faceid_sub_set = set()
    file_count = 0
    with open(face_list_path, 'r') as file:
        for line in file.readlines():
            file_count += 1
            if file_count % 10000 == 0:
                print('Read File Count: {}'.format(file_count))

            # if file_count == 100000:
            #     break

            # 015ff272-3768-4d4c-bebc-c14dc93065cd 477fc604-1479-4745-b8e9-9440a6ac8187 f8e4735b44724e25077106e0eaef15f0.jpg 0
            lines = line.strip().split('_')

            faceid = lines[0]
            sub_faceid = lines[1]

            faceid_sub = '{} {}'.format(faceid, sub_faceid)
            faceid_sub_set.add(faceid_sub)
        print('Read File Count: {}'.format(file_count))
    return faceid_sub_set


def read_merge_list(face_list_path):
    faceid_sub_set = set()
    file_count = 0
    with open(face_list_path, 'r') as file:
        for line in file.readlines():
            file_count += 1
            if file_count % 10000 == 0:
                print('Read File Count: {}'.format(file_count))

            # if file_count == 100000:
            #     break

            # 015ff272-3768-4d4c-bebc-c14dc93065cd 477fc604-1479-4745-b8e9-9440a6ac8187 f8e4735b44724e25077106e0eaef15f0.jpg 0
            lines = line.strip().split('_')

            faceid = lines[0]
            sub_faceid = lines[1]

            faceid_sub = '{} {}'.format(faceid, sub_faceid)
            faceid_sub_set.add(faceid_sub)
        print('Read File Count: {}'.format(file_count))
    return faceid_sub_set


def read_merge_file(file_list):
    faceid_subfaceids = []
    file_count = 0
    with open(file_list, 'r') as file:
        for line in file.readlines():
            file_count += 1
            if file_count % 10000 == 0:
                print('Read Feature Count: {}'.format(file_count))

            # if file_count == 10000:
            #     break

            # NONE:c651b1e5-fdbc-4447-9a6a-9213d7729faa_ac638be4-0e9a-42e8-a29b-c730a41aac50,752524c4-e540-45e2-8388-55aab4d72150_94d61342-0841-432a-a8f2-5bb6cf5c9a4f
            lines = line.strip().split(':')
            # xiaoqu = lines[0]
            faceids = lines[1].split(',')
            faceids = [faceid.replace('_', ' ', 1) for faceid in faceids]

            faceid_subfaceids.append(faceids)
        print('Read Merge Count: {}'.format(file_count))
    return faceid_subfaceids


def get_most_image_count_faceids(faceid_subfaceids_list, faceid_subfaceid_file_feature_names):
    faceid_sub_set = set()
    for faceid_subfaceids in faceid_subfaceids_list:
        counts = []
        for faceid_subfaceid in faceid_subfaceids:
            file_feature_names = faceid_subfaceid_file_feature_names.get(faceid_subfaceid, None)
            if file_feature_names is None:
                counts.append(0)
                continue

            file_infos, features, file_names = file_feature_names
            image_count = len(file_names)
            counts.append(image_count)
        counts = np.array(counts, dtype=np.int32)
        idx = counts.argmax()
        if counts[idx] > 1:
            faceid_sub_set.add(faceid_subfaceids[idx])

    return faceid_sub_set


def main():
    parser = argparse.ArgumentParser(description="Select Face by Attr")
    parser.add_argument(
        "--file_list", default="/rootfs/media/kasim/Data1/data/VideoFaceImage/filter6/FaceRecSel.txt", help="path to all face image file list", type=str
    )
    parser.add_argument(
        "--feature_file", default="/rootfs/media/kasim/Data1/data/VideoFaceImage/filter6/FaceRecFeatSel.dat", help="path to feature file", type=str
    )
    parser.add_argument(
        "--face_list", default="/rootfs/media/kasim/Data1/data/VideoFaceImage/filter6/DiffFace.0.4.txt", help="path to face id list", type=str
    )
    parser.add_argument(
        "--merge_list", default="/rootfs/media/kasim/Data1/data/VideoFaceImage/filter6/MergeFace.0.4.txt", help="path to merge faceid list", type=str
    )
    parser.add_argument(
        "--attr_file", default="/rootfs/media/kasim/Data1/data/VideoFaceImage/filter3/FaceAttrs.txt", help="path to merge file", type=str
    )
    parser.add_argument(
        "--output_dir", default="/rootfs/media/kasim/Data1/data/VideoFaceImage/filter6", help="path to output", type=str
    )
    parser.add_argument(
        "--feature_dim", default=256, help="feature dim", type=int
    )
    parser.add_argument('--min_face_image_count', type=int, default=20, help='min face image count')
    parser.add_argument('--mode', type=int, default=1, help='mode')
    args = parser.parse_args()

    output_dir = args.output_dir
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        os.system('chmod a+wr {}'.format(output_dir))

    faceid_subfaceid_file_feature_names = read_features(args.file_list, args.feature_file, args.feature_dim)
    name_attrs = read_attr_file(args.attr_file)
    faceid_sub_set0 = read_face_list(args.face_list)
    faceid_subfaceids_list = read_merge_file(args.merge_list)

    out_file_list_path = os.path.join(output_dir, 'AttrFaceRecSel.txt')
    out_feature_file_path = os.path.join(output_dir, 'AttrFaceRecFeatSel.dat')
    out_file_list_file = open(out_file_list_path, 'w')
    out_feature_file = open(out_feature_file_path, 'wb')

    out_face_list_path = os.path.join(output_dir, 'AttrDiffFace.txt')
    out_image_list_path = os.path.join(output_dir, 'AttrDiffFaceImage.txt')
    out_face_list_file = open(out_face_list_path, 'w')
    out_image_list_file = open(out_image_list_path, 'w')

    face_count = 0

    face_image_count = 0

    faceid_sub_set1 = get_most_image_count_faceids(faceid_subfaceids_list, faceid_subfaceid_file_feature_names)

    for faceid_sub_set in [faceid_sub_set0, faceid_sub_set1]:
        for faceid_subfaceid in faceid_sub_set:
            file_feature_names = faceid_subfaceid_file_feature_names.get(faceid_subfaceid, None)
            if file_feature_names is None:
                continue
            file_infos, features, file_names = file_feature_names
            image_count = len(file_names)
            if image_count < args.min_face_image_count:
                continue

            attrs = [name_attrs[name] for name in file_names]
            pitchs = []
            rolls = []
            yaws = []
            qualitys = []
            for [pitch, roll, yaw], quality in attrs:
                pitchs.append(pitch)
                rolls.append(roll)
                yaws.append(yaw)
                qualitys.append(quality)
            pitchs = np.array(pitchs, dtype=np.float32)
            rolls = np.array(rolls, dtype=np.float32)
            yaws = np.array(yaws, dtype=np.float32)
            qualitys = np.array(qualitys, dtype=np.float32)

            idxes = np.arange(0, image_count)
            pitchs_mask = np.logical_and(MIN_PITCH1 <= pitchs, pitchs <= MAX_PITCH1)
            rolls_mask = np.logical_and(MIN_ROLL1 <= rolls, rolls <= MAX_ROLL1)
            yaws_mask = np.logical_and(MIN_YAW1 <= yaws, yaws <= MAX_YAW1)
            mask = pitchs_mask & rolls_mask & yaws_mask
            if mask.any():
                max_quality_idx = qualitys[mask].argmax()
                idx = idxes[mask][max_quality_idx]
            else:
                pitchs_mask = np.logical_and(MIN_PITCH2 <= pitchs, pitchs <= MAX_PITCH2)
                rolls_mask = np.logical_and(MIN_ROLL2 <= rolls, rolls <= MAX_ROLL2)
                yaws_mask = np.logical_and(MIN_YAW2 <= yaws, yaws <= MAX_YAW2)
                mask = pitchs_mask & rolls_mask & yaws_mask
                if mask.any():
                    max_quality_idx = qualitys[mask].argmax()
                    idx = idxes[mask][max_quality_idx]
                else:
                    if args.mode == 0:
                        continue
                    else:
                        idx = qualitys.argmax()

            face_count += 1
            face_image_count += image_count

            for file_name, feature in zip(file_names, features):
                file_info = '{} {} {}\n'.format(faceid_subfaceid, file_name, image_count)
                out_file_list_file.write(file_info)
                out_feature_file.write(feature)
            faceid_info = faceid_subfaceid.replace(' ', '_', 1) + '\n'
            out_face_list_file.write(faceid_info)
            out_image_list_file.write(file_names[idx] + '\n')
    out_feature_file.close()
    out_file_list_file.close()
    out_face_list_file.close()
    out_image_list_file.close()
    os.system('chmod a+wr {}'.format(out_face_list_path))
    os.system('chmod a+wr {}'.format(out_image_list_path))
    os.system('chmod a+wr {}'.format(out_file_list_path))
    os.system('chmod a+wr {}'.format(out_feature_file_path))
    print('Face Count: {}, Face Image Count: {}'.format(face_count, face_image_count))
    print('finish!')


if __name__ == '__main__':
    main()

